package com.alibaba.ttl.threadpool;

import cn.micro.ttl.mdc.MDCThreadWrap;
import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.spi.TtlEnhanced;
import com.alibaba.ttl.spi.TtlWrapper;
import org.springframework.core.task.AsyncTaskExecutor;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;

public class AsyncTaskExecutorTtlMDCWrapper implements TtlWrapper<AsyncTaskExecutor>, TtlEnhanced, AsyncTaskExecutor {

    private final AsyncTaskExecutor asyncTaskExecutor;
    protected final boolean idempotent;

    public AsyncTaskExecutorTtlMDCWrapper(AsyncTaskExecutor asyncTaskExecutor, boolean idempotent) {
        this.asyncTaskExecutor = asyncTaskExecutor;
        this.idempotent = idempotent;
    }

    @Override
    public void execute(Runnable command) {
        asyncTaskExecutor.execute(TtlRunnable.get(MDCThreadWrap.runnable(command), false, idempotent));
    }

    @Override
    public AsyncTaskExecutor unwrap() {
        return asyncTaskExecutor;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;

        AsyncTaskExecutorTtlMDCWrapper that = (AsyncTaskExecutorTtlMDCWrapper) o;

        return asyncTaskExecutor.equals(that.asyncTaskExecutor);
    }

    @Override
    public int hashCode() {
        return asyncTaskExecutor.hashCode();
    }

    @Override
    public String toString() {
        return this.getClass().getName() + " - " + asyncTaskExecutor.toString();
    }


    @Override
    public void execute(Runnable task, long startTimeout) {
        asyncTaskExecutor.execute(TtlRunnable.get(MDCThreadWrap.runnable(task), false, idempotent), startTimeout);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return asyncTaskExecutor.submit(TtlRunnable.get(MDCThreadWrap.runnable(task), false, idempotent));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return asyncTaskExecutor.submit(TtlCallable.get(MDCThreadWrap.callable(task), false, idempotent));
    }
}
